import ast
import re
import openpyxl
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment, Border, Side, PatternFill
from openpyxl.utils import get_column_letter


def extract_all_classes_fields(file_path):
    """
    从Python文件中提取所有类的字段信息

    参数:
    file_path: Python文件路径

    返回:
    字典，键为类名，值为包含字段名和中文名的列表，格式为{class_name: [(field_name, chinese_name), ...]}
    """
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()
    
    # 解析Python代码为AST
    tree = ast.parse(content)
    
    classes_fields = {}
    
    # 遍历AST节点
    for node in ast.walk(tree):
        # 查找类定义
        if isinstance(node, ast.ClassDef):
            class_name = node.name
            fields = []
            
            # 遍历类中的节点
            for class_node in node.body:
                # 查找赋值语句（字段定义）
                if isinstance(class_node, ast.Assign):
                    for target in class_node.targets:
                        if isinstance(target, ast.Name):
                            field_name = target.id
                            chinese_name = extract_chinese_name(class_node)
                            fields.append((field_name, chinese_name))
            
            # 只添加有字段的类
            if fields:
                classes_fields[class_name] = fields
    
    return classes_fields


def extract_chinese_name(assign_node):
    """
    从赋值节点中提取中文名

    参数:
    assign_node: AST赋值节点

    返回:
    字段的中文名，如果没有找到则返回空字符串
    """
    # 检查赋值语句的值是否是函数调用（如字段定义）
    if isinstance(assign_node.value, ast.Call):
        call = assign_node.value
        # 遍历关键字参数
        for keyword in call.keywords:
            if keyword.arg == 'verbose_name':
                # 如果是字符串字面量
                if isinstance(keyword.value, ast.Str):
                    return keyword.value.s
                # 如果是常量（Python 3.8+）
                elif hasattr(ast, 'Constant') and isinstance(keyword.value, ast.Constant):
                    return keyword.value.value
    
    return ""


def write_to_excel(classes_fields, output_file):
    """
    将字段信息写入Excel文件，每个类一个工作表

    参数:
    classes_fields: 包含类名和字段信息的字典
    output_file: 输出的Excel文件路径
    """
    # 创建新的工作簿
    wb = Workbook()
    
    # 删除默认创建的工作表
    if 'Sheet' in wb.sheetnames:
        wb.remove(wb['Sheet'])
    
    # 定义样式
    header_font = Font(bold=True, color="FFFFFF")
    header_fill = PatternFill(start_color="366092", end_color="366092", fill_type="solid")
    header_alignment = Alignment(horizontal="center", vertical="center")
    border = Border(
        left=Side(style='thin'),
        right=Side(style='thin'),
        top=Side(style='thin'),
        bottom=Side(style='thin')
    )
    
    # 为每个类创建一个工作表
    for class_name, fields in classes_fields.items():
        # 创建工作表，工作表名称不能超过31个字符且不能包含特殊字符
        sheet_name = class_name[:31]  # Excel限制工作表名称长度
        sheet_name = re.sub(r'[\\/*?[\]:]', '', sheet_name)  # 移除非法字符
        
        # 如果工作表名称已存在，添加数字后缀
        original_name = sheet_name
        counter = 1
        while sheet_name in wb.sheetnames:
            sheet_name = f"{original_name}_{counter}"
            counter += 1
        
        ws = wb.create_sheet(sheet_name)
        
        # 添加表头（从B2单元格开始）
        # ws['B2'] = "字段名"
        # ws['B3'] = "中文名"
        
        # 应用表头样式
        for row in [2, 3]:  # 第2行和第3行
            cell = ws[f'B{row}']
            cell.font = header_font
            cell.fill = header_fill
            cell.alignment = header_alignment
            cell.border = border
        
        # 写入数据（从C2和C3单元格开始）
        for i, (field_name, chinese_name) in enumerate(fields, start=2):
            col_letter = get_column_letter(i + 1)  # 从C列开始 (i=2 -> C列)
            
            # 写入英文名（第2行）
            ws[f'{col_letter}2'] = field_name
            cell = ws[f'{col_letter}2']
            cell.border = border
            
            # 写入中文名（第3行）
            ws[f'{col_letter}3'] = chinese_name
            cell = ws[f'{col_letter}3']
            cell.border = border
        
        # 添加5行空行（从第4行开始）
        for row in range(4, 9):  # 第4行到第8行
            for col in range(2, len(fields) + 3):  # 从B列到最后一个字段列
                col_letter = get_column_letter(col)
                cell = ws[f'{col_letter}{row}']
                cell.border = border
        
        # 自动调整列宽
        adjust_column_width(ws)
    
    # 保存文件
    wb.save(output_file)
    print(f"数据已写入 {output_file}，共处理了 {len(classes_fields)} 个类")


def adjust_column_width(ws):
    """
    自动调整列宽以适应内容

    参数:
    ws: Excel工作表对象
    """
    for column in ws.columns:
        max_length = 0
        column_letter = get_column_letter(column[0].column)
        for cell in column:
            try:
                if cell.value and len(str(cell.value)) > max_length:
                    max_length = len(str(cell.value))
            except:
                pass
        adjusted_width = min(max_length + 2, 50)  # 限制最大宽度为50
        ws.column_dimensions[column_letter].width = adjusted_width


# 使用示例
if __name__ == "__main__":
    # 替换为你的Python文件路径
    file_path = "../house/models.py"
    # 输出的Excel文件路径
    output_file = "all_model_fields.xlsx"
    
    # 提取所有类的字段信息
    classes_fields = extract_all_classes_fields(file_path)
    
    # 如果没有找到任何类，打印警告
    if not classes_fields:
        print("警告：没有找到任何包含字段的类")
    else:
        # 打印找到的类
        print(f"找到 {len(classes_fields)} 个类:")
        for class_name in classes_fields.keys():
            print(f"  - {class_name}")
        
        # 写入Excel
        write_to_excel(classes_fields, output_file)